# load image and preprocess it with vgg16 structure
# --by flare
import numpy as np
from keras.applications.vgg16 import VGG16
from keras.applications.vgg16 import preprocess_input
from keras.preprocessing.image import img_to_array, load_img

model_vgg = VGG16(weights='imagenet', include_top=False)


# define a method to load and preprocess the image
def modelProcess(img_path, model):
    img = load_img(img_path, target_size=(224, 224))
    img = img_to_array(img)
    x = np.expand_dims(img, axis=0)
    x = preprocess_input(x)
    x_vgg = model.predict(x)
    x_vgg = x_vgg.reshape(1, 25088)
    return x_vgg


# list file names of the training datasets
import os

folder = "dataset/data_vgg/cats"
dirs = os.listdir(folder)
# generate path for the images
img_path = []
for i in dirs:
    if os.path.splitext(i)[1] == ".jpg":
        img_path.append(i)
img_path = [folder + "//" + i for i in img_path]

# preprocess multiple images
features1 = np.zeros([len(img_path), 25088])
for i in range(len(img_path)):
    feature_i = modelProcess(img_path[i], model_vgg)
    print('preprocessed:', img_path[i])
    features1[i] = feature_i

folder = "dataset/data_vgg/dogs"
dirs = os.listdir(folder)
img_path = []
for i in dirs:
    if os.path.splitext(i)[1] == ".jpg":
        img_path.append(i)
img_path = [folder + "//" + i for i in img_path]
features2 = np.zeros([len(img_path), 25088])
for i in range(len(img_path)):
    feature_i = modelProcess(img_path[i], model_vgg)
    print('preprocessed:', img_path[i])
    features2[i] = feature_i

# label the results
print(features1.shape, features2.shape)
y1 = np.zeros(300)
y2 = np.ones(300)

# generate the training 笔记.md
X = np.concatenate((features1, features2), axis=0)
y = np.concatenate((y1, y2), axis=0)
y = y.reshape(-1, 1)
print(X.shape, y.shape)

# 数据分离
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=50)
print(X_train.shape, y_train.shape)
print(X_test.shape, y_test.shape)

# 建立模型
from keras.models import Sequential
from keras.layers import Dense

model = Sequential()
model.add(Dense(units=10, activation='relu', input_dim=25088))
model.add(Dense(units=1, activation='sigmoid'))
model.summary()
# 配置模型
model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
# 训练模型
model.fit(X_train, y_train, epochs=50)
# 测试模型
score = model.evaluate(X_test, y_test, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])

# 预测新图片
img_path = 'pic_dog.jpeg'
img = load_img(img_path, target_size=(224, 224))
img = img_to_array(img)
model_vgg = VGG16(weights='imagenet', include_top=False)
x = np.expand_dims(img, axis=0)
x = preprocess_input(x)
# 特征提取
feature_vgg = model_vgg.predict(x)
feature_vgg = feature_vgg.reshape(1, 7 * 7 * 512)
result = model.predict(feature_vgg)
# print(X_train.class_indices)
print('result：', result)

# 预测新图片
img_path = 'pic_cat.jpeg'
img = load_img(img_path, target_size=(224, 224))
img = img_to_array(img)
model_vgg = VGG16(weights='imagenet', include_top=False)
x = np.expand_dims(img, axis=0)
x = preprocess_input(x)
# 特征提取
feature_vgg = model_vgg.predict(x)
feature_vgg = feature_vgg.reshape(1, 7 * 7 * 512)
result = model.predict(feature_vgg)
# print(X_train.class_indices)
print('result：', result)